package org.sun.basic.util;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

import junit.framework.Assert;

public class AbstractDbUnitTestCase {
	public static IDatabaseConnection dbunitCon;
	private File tempFile;
	
	@BeforeClass
	public static void init() throws DatabaseUnitException, SQLException{
		dbunitCon= new DatabaseConnection(DbUtil.getConnection());
	}
	
	protected IDataSet createDataSet(String tname) throws DataSetException{
		InputStream is= AbstractDbUnitTestCase.class
				.getClassLoader().getResourceAsStream(tname+".xml");
		Assert.assertNotNull("dbunit文件不能为空！", is);
		return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is)));
	}

	protected void backupAllTable() throws SQLException, DataSetException, IOException{
		IDataSet ds=dbunitCon.createDataSet();
		writeBackupFile(ds);
	}
	
	protected void writeBackupFile(IDataSet ds) throws SQLException, IOException, DataSetException{
		tempFile=File.createTempFile("backup", "xml");
		FlatXmlDataSet.write(ds, new FileWriter(tempFile));
	}
	
	protected void backupCustomTable(String[] tables) throws DataSetException, SQLException, IOException{
		QueryDataSet ds= new QueryDataSet(dbunitCon);
		for(String table:tables){
			ds.addTable(table);
		}
		writeBackupFile(ds);
	}
	
	protected void backupOneTable(String table) throws DataSetException, SQLException, IOException{
		this.backupCustomTable(new String[] {table});
	}
	
	protected void resumeTable() throws FileNotFoundException, DatabaseUnitException, SQLException{
		IDataSet ds= new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tempFile))));
		DatabaseOperation.CLEAN_INSERT.execute(dbunitCon, ds);
	}
	@AfterClass
	public static void destory(){
		if(dbunitCon!=null)
			try {
				dbunitCon.close();
			} catch (SQLException e) {
				e.printStackTrace();
			}
	}
}
